'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2025-08-06 20:35:21
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

import networkx as nx
import json
import pickle


def tag_to_graph(tag_list: list[str]):
    G = nx.DiGraph()
    for tag_str in tag_list:
        if '/' not in tag_str:
            G.add_node(tag_str)
        else:
            for edge in tag_str_to_edges(tag_str):
                G.add_edge(edge[0], edge[1])
    return G


def tag_str_to_edges(tag_str):
    ''' tag1/tag2/tag3 -> [(tag1, tag1/tag2), (tag1/tag2, tag1/tag2/tag3)] '''
    # processing tags
    tag_list = tag_str.split('/')
    for i in range(len(tag_list) - 1):
        yield '/'.join(tag_list[:i + 1]), '/'.join(tag_list[:i + 2])


class LiteratureViewer:
    literature_list: list[dict]
    selected_literature: list[int]  # indices in self.literature_list
    kw_graph: nx.DiGraph
    selected_kw: set[str]
    hidden_items: set[int]
    sort_key: tuple[str, bool]

    def __init__(self) -> None:
        self.load_data([])
    
    def load_json(self, fname: str):
        with open(fname, 'r', encoding='utf-8') as f:
            literature_list = json.load(f)
        self.load_data(literature_list)
    
    def load_pickle(self, fname: str):
        with open(fname, 'rb') as f:
            literature_list = pickle.load(f)
        self.load_data(literature_list)
    
    def load_from_file(self, fname: str):
        if fname.endswith('.json'):
            self.load_json(fname)
        elif fname.endswith('.pickle'):
            self.load_pickle(fname)
        else:
            raise ValueError(f'Unknown file type: {fname}')

    def load_data(self, literature_list: list[dict]):
        self.literature_list = literature_list
        self.selected_literature = list(range(len(literature_list)))
        self.kw_graph = self.get_kw_from_selected_items()
        self.selected_kw = set()
        self.hidden_items = set()
        self.sort_by('title')

    def get_kw_from_selected_items(self):
        all_kw = set()
        for i in self.selected_literature:
            all_kw.update(self.literature_list[i]['keywords'])
        return tag_to_graph(all_kw)

    def select_item_from_kw(self):
        self.selected_literature = []
        for item_idx, item in enumerate(self.literature_list):
            if item_idx in self.hidden_items:
                continue

            curr_flag = True  # whether the item is selected
            for target_kw in self.selected_kw:
                curr_flag = False
                for item_kw in item["keywords"]:
                    if item_kw.startswith(target_kw):
                        curr_flag = True
                        break
                if not curr_flag:
                    break
            if curr_flag:
                self.selected_literature.append(item_idx)

    def sort_by(self, key_name: str, reverse=False):
        self.sort_key = (key_name, reverse)

    def select_kw(self, kw: str):
        self.selected_kw.add(kw)
        self.select_item_from_kw()
        self.kw_graph = self.get_kw_from_selected_items()
    
    def deselect_kw(self, kw: str):
        self.selected_kw.remove(kw)
        self.select_item_from_kw()
        self.kw_graph = self.get_kw_from_selected_items()

    def get_kw_graph(self):
        ''' Print the keyword graph by nested dicts.
            Example: for graph
                tag1 -> tag2 -> tag3
                tag4 -> tag5 -> tag6
            the output should be:
                {
                    "tag1": {
                        "tag2": {
                            "tag3": {}
                        }
                    },
                    "tag4": {
                        "tag5": {
                            "tag6": {}
                        }
                    }
                }
            If the graph is empty, return empty dict.
        '''
        if self.kw_graph.number_of_nodes() == 0:
            return {}
        else:
            return nx.to_dict_of_dicts(self.kw_graph)

    def get_kw_leaf(self):
        ''' Get all leaf nodes of self.kw_graph, return as a list. '''
        return [node for node in self.kw_graph.nodes if self.kw_graph.out_degree(node) == 0]

    def hide_item(self, item_idx: int):
        self.hidden_items.add(item_idx)
        self.select_item_from_kw()
        self.kw_graph = self.get_kw_from_selected_items()

    def get_selected_keywords(self):
        return self.selected_kw

    def get_all(self):
        # 创建排序后的文献索引列表
        sorted_literature = sorted(self.selected_literature, 
                                 key=lambda x: self.literature_list[x][self.sort_key[0]], 
                                 reverse=self.sort_key[1])
        
        all_info = {
            "selected_idx": sorted_literature,
            "sort_key": self.sort_key,
            "hidden_idx": list(self.hidden_items),
            "kw_graph": self.get_kw_leaf(),
            "selected_kw": list(self.selected_kw)
        }
        # print(all_info)
        return all_info

    def recover_item(self, item_idx: int):
        self.hidden_items.remove(item_idx)
        self.select_item_from_kw()
        self.kw_graph = self.get_kw_from_selected_items()


if __name__ == "__main__":
    tag_list = ["tag1/tag2/tag3", "tag4"]
    print(tag_str_to_edges("tag1/tag2/tag3"))
    G = tag_to_graph(tag_list)
    print(G.nodes)
    print(G.edges)
